﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Net.Sockets;
using System.Net;

namespace Alex.Library.Net.Communications
{
    /// <summary>
    /// Tcp数据分发器
    /// </summary>
    public class TcpDispatcher : IDisposable
    {
        /// <summary>
        /// 默认构造函数
        /// </summary>
        public TcpDispatcher()
        {
            _listeners = new Dictionary<string, TcpListener>();
        }

        private Dictionary<string, TcpListener> _listeners;

        /// <summary>
        /// 启动监听
        /// </summary>
        public void Start(string id, ListeningArg arg)
        {
            try
            {
                if (_listeners.ContainsKey(id))
                {
                    throw new Exception("监听标识已存在.");
                }
                if (arg == null)
                {
                    throw new NullReferenceException("监听参数不能为null.");
                }
                if (arg.LocalPort < 1000)
                {
                    throw new Exception("监听的本地端口号不能小于1000.");
                }
                if (arg.RemoteEP == null)
                {
                    throw new NullReferenceException("客户端IP终结点不能为null.");
                }

                TcpListener _listener = new TcpListener(IPAddress.Any, arg.LocalPort);
                _listener.Start();
                _listener.BeginAcceptSocket(new AsyncCallback(OnSocketAccept), new object[] { _listener, arg.RemoteEP });
                _listeners[id] = _listener;
            }
            catch
            {
                throw;
            }
        }

        /// <summary>
        /// 关闭指定标识的监听
        /// </summary>
        /// <param name="id">监听标识</param>
        public void Stop(string id)
        {
            try
            {
                if (_listeners.ContainsKey(id))
                {
                    _listeners[id].Stop();
                }
            }
            catch
            {
            }
        }

        /// <summary>
        /// 关闭全部监听
        /// </summary>
        public void StopAll()
        {
            try
            {
                foreach (TcpListener listener in _listeners.Values)
                {
                    listener.Stop();
                }
            }
            catch
            {
            }
        }

        private void OnSocketAccept(IAsyncResult iar)
        {
            object[] userstate = iar.AsyncState as object[];
            TcpListener listener = userstate[0] as TcpListener;
            IPEndPoint remoteEP = userstate[1] as IPEndPoint;
            Socket sin_client = listener.EndAcceptSocket(iar);
            DispatchSession ds = new DispatchSession();
            ds.Dispatching(sin_client, remoteEP);
            listener.BeginAcceptSocket(new AsyncCallback(OnSocketAccept), userstate);
        }

        public void Dispose()
        {
            StopAll();
            foreach (string id in _listeners.Keys)
            {
                _listeners[id] = null;
            }
            _listeners.Clear();
        }
    }

    public class ListeningArg
    {        
        public int LocalPort { get; set; }
        public IPEndPoint RemoteEP { get; set; }
    }

    class DispatchSession
    {
        public DispatchSession()
        {

        }

        private Socket _sin_client;
        private Socket _sin_server;
        private IPEndPoint _forwardEP;

        /// <summary>
        /// 分发数据
        /// </summary>
        /// <param name="sin_client"></param>
        /// <param name="forwardEP"></param>
        public void Dispatching(Socket sin_client, IPEndPoint forwardEP)
        {
            this._sin_client = sin_client;
            this._forwardEP = forwardEP;

            SocketAsyncEventArgs client_args = new SocketAsyncEventArgs();
            client_args.SetBuffer(new byte[1024], 0, 1024);
            client_args.RemoteEndPoint = _sin_client.RemoteEndPoint;
            client_args.UserToken = _sin_client;
            client_args.Completed += client_args_Completed;
            _sin_client.ReceiveAsync(client_args);
        }

        private void client_args_Completed(object sender, SocketAsyncEventArgs e)
        {
            if (e.BytesTransferred > 0)
            {
                try
                {
                    if (_sin_server == null)
                    {
                        _sin_server = new TcpClient().Client;
                        _sin_server.Connect(_forwardEP);

                        SocketAsyncEventArgs svcrearg = new SocketAsyncEventArgs();
                        svcrearg.SetBuffer(new byte[1024], 0, 1024);
                        svcrearg.RemoteEndPoint = _sin_server.RemoteEndPoint;
                        svcrearg.UserToken = _sin_server;
                        svcrearg.Completed += server_args_Completed;
                        _sin_server.ReceiveAsync(svcrearg);
                    }

                    Forwarding(_sin_client, _sin_server, e);
                }
                catch
                {
                    throw;
                }
            }
        }

        private void server_args_Completed(object sender, SocketAsyncEventArgs e)
        {
            if (e.BytesTransferred > 0)
            {
                Forwarding(_sin_server, _sin_client, e);
            }
        }

        private void Forwarding(Socket src, Socket dest, SocketAsyncEventArgs e)
        {
            if (e.BytesTransferred == 0)
            {
                return;
            }
            if (!ValidateState(dest, SelectMode.SelectWrite))
            {
                return;
            }
            SocketError errorCode = SocketError.Success;
            dest.Send(e.Buffer, e.Offset, e.BytesTransferred, SocketFlags.None, out errorCode);
            if (errorCode != SocketError.Success)
            {
                CloseAll();
            }
            try
            {
                src.ReceiveAsync(e);
            }
            catch
            {
                CloseAll();
            }
        }

        private bool ValidateState(Socket socket, params SelectMode[] selectModes)
        {
            if (!socket.Connected)
            {
                CloseAll();
                return false;
            }
            try
            {
                foreach (var selectMode in selectModes)
                {
                    var ret = socket.Poll(1000 * 500, selectMode);
                    if (!ret)
                    {
                        CloseAll();
                        return false;
                    }
                }
            }
            catch
            {
                CloseAll();
                return false;
            }
            return true;
        }

        private void CloseAll()
        {
            try
            {
                _sin_client.Close();
            }
            catch
            {
            }

            try
            {
                if (_sin_server != null)
                {
                    _sin_server.Close();
                }
            }
            catch
            {
            }
        }
    }
}
